import json
import os

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

plt.rcParams['xtick.direction'] = 'in'
plt.rcParams['ytick.direction'] = 'in'
plt.rcParams['axes.titlesize'] = 24
plt.rcParams['font.size'] = 24
data_dir = os.path.join(os.getcwd(), '..', 'experiment')


def plot_distribution_kmeans():
    with open(os.path.join(data_dir, 'data_01_02.json'), 'r') as f:
        kmeans_data = json.load(f)
    fig, ax = plt.subplots(figsize=(9, 8), dpi=300)
    for i in range(len(kmeans_data)):
        part = kmeans_data[i]
        x = np.array(part['ue'])[:, 0].squeeze()
        y = np.array(part['ue'])[:, 1].squeeze()
        ax.scatter(x=x, y=y, marker='o', c='#000', s=80)
        center = part['center']
        draw_circle = plt.Circle((center[0], center[1]), radius=220, facecolor="None", edgecolor="#000", linestyle=':', linewidth=4)
        ax.add_artist(draw_circle)
    line_x = [76.070, 443.579, 467.986, 755.601, 888.714, 722.411, 495.616, 158.192]
    line_y = [712.596, 684.933, 907.351, 821.316, 455.136, 594.659, 507.722, 174.355]
    line = plt.Line2D(line_x, line_y, color='#1B7C8C', linestyle='-', markersize=20, linewidth=8, label="Trajectory (85.63s)")
    line.set_markeredgewidth(2)
    ax.add_line(line)
    ax.scatter(x=line_x[0], y=line_y[0]+5, marker='^', edgecolor='#1B7C8C', s=600, label="Start point", facecolor="None", linewidths=4)
    ax.scatter(x=line_x[-1], y=line_y[-1], marker='s', edgecolor='#1B7C8C', s=600, label="End point", facecolor="None", linewidths=4)
    ax.scatter(x=line_x[1:-1], y=line_y[1:-1], marker='o', edgecolor='#1B7C8C', s=600, facecolor="None", linewidths=4)
    plt.legend()
    ax.set_xlabel('X(m)', loc='right')
    ax.set_ylabel('Y(m)', loc='top')
    ax.axis([0, 1000, 1, 1000])
    ax.grid(True, color="#999", linestyle='--')
    plt.tight_layout()
    plt.savefig(os.path.join(os.path.expanduser("~/Desktop"), "ex_01.png"))


def plot_distribution_wusc():
    with open(os.path.join(data_dir, 'data_01_01.json'), 'r') as f:
        welzl_data = json.load(f)
    fig, ax = plt.subplots(figsize=(9, 8), dpi=300)
    for i in range(len(welzl_data)):
        part = welzl_data[i]
        x = np.array(part['ue'])[:, 0].squeeze()
        y = np.array(part['ue'])[:, 1].squeeze()
        ax.scatter(x=x, y=y, marker='o', c='#000', s=80)
        center = part['center']
        draw_circle = plt.Circle((center[0], center[1]), radius=220, facecolor="None", edgecolor="#000", linestyle=':', linewidth=4)
        ax.add_artist(draw_circle)
    line_x = [159.890, 523.736, 512.941, 824.446, 815.363, 827.973, 131.839]
    line_y = [734.951, 576.064, 861.294, 831.232, 540.687, 399.138, 181.725]
    line = plt.Line2D(line_x, line_y, color='#1B7C8C', linestyle='-', linewidth=8, label="Trajectory (71.83s)")
    ax.add_line(line)
    ax.scatter(x=line_x[0], y=line_y[0]+10, marker='^', edgecolor='#1B7C8C', s=600, label="Start point", facecolor="None", linewidths=4)
    ax.scatter(x=line_x[-1], y=line_y[-1], marker='s', edgecolor='#1B7C8C', s=600, label="End point", facecolor="None", linewidths=4)
    ax.scatter(x=line_x[1:-1], y=line_y[1:-1], marker='o', edgecolor='#1B7C8C', s=600, facecolor="None", linewidths=4)
    plt.legend()
    ax.set_xlabel('X(m)', loc='right')
    ax.set_ylabel('Y(m)', loc='top')
    ax.axis([0, 1000, 1, 1000])
    ax.grid(True, color="#999", linestyle='--')
    plt.tight_layout()
    plt.savefig(os.path.join(os.path.expanduser("~/Desktop"), "ex_02.png"))


def plot_partition():
    file = os.path.join(data_dir, 'data_01_03.csv')
    df = pd.read_csv(file, header=0, delimiter=' ')
    data = df.values
    plt.figure(figsize=(16, 8), dpi=300)
    bar_width = 0.45
    x = data[:, 0] / 1000
    y1 = data[:, 1]
    y2 = data[:, 2]

    plt.bar(x, y1, color='#2F4858', width=bar_width, label='KMeans')
    plt.bar(x + bar_width, y2, color='#46D0AA', width=bar_width, label='WUSC')
    plt.plot(x, y1+10, marker='x', markersize=12, linestyle='-', color='#2F4858', linewidth=3)
    plt.plot(x + bar_width, y2+10, marker='x', markersize=12, linestyle='-', color='#46D0AA', linewidth=3)
    plt.axhline(y=1034, color='#874C4A', linestyle='--', label='Max', linewidth=3)
    plt.xlabel('Number of UDs (k)')
    plt.ylabel('Number of PTs')
    plt.grid(True, linestyle=':', color='#999')
    plt.legend(loc='upper left')
    plt.savefig(os.path.join(os.path.expanduser("~/Desktop"), "ex_03.png"))


def plot_sr():
    data1 = pd.read_csv(os.path.join(data_dir, 'data_02_01.csv')).values
    data2 = pd.read_csv(os.path.join(data_dir, 'data_02_02.csv')).values
    data3 = pd.read_csv(os.path.join(data_dir, 'data_02_03.csv')).values
    data4 = pd.read_csv(os.path.join(data_dir, 'data_02_08.csv'), header=None).values
    plt.figure(figsize=(9, 8), dpi=300)
    y1 = data1[:100, 0][1::5]/100
    y2 = data2[:1000, 0][1::50]/100
    y3 = data3[:100, 0][1::5] / 100
    y1 = y1.tolist()
    y1.insert(0, 1)
    y2 = y2.tolist()
    y2.insert(0, 1)
    y3 = y3.tolist()
    y3.insert(0, 1)
    x = [i*5 for i in range(len(y1))]
    plt.plot(x, y1, marker='^', markerfacecolor='none', markersize=16, markeredgewidth=3, linestyle=':', color='#2F4858', label='CD-Greedy', linewidth=5)
    plt.plot(x, y3, marker='x', markerfacecolor='none', markersize=16, markeredgewidth=3, linestyle=':', color='#19B4A8', label='CD-OPT', linewidth=5)
    plt.annotate('It=100', xy=(x[-1], y3[-1]), xytext=(x[-1]-30, y3[-1]-0.028), arrowprops=dict(facecolor='#874C4A', arrowstyle='->'))
    plt.annotate('  ', xy=(x[-1], y1[-1]), xytext=(x[-1]-10, y1[-1]+0.015), arrowprops=dict(facecolor='#874C4A', arrowstyle='->'))
    x4 = data4[:, 0]
    y4 = data4[:, 2] / 100
    plt.plot(x4, y4, marker='s', markerfacecolor='none', markersize=16, markeredgewidth=3, linestyle=':', color='#874C4A', label='Random', linewidth=5)
    plt.annotate('It=50000', xy=(x4[-1], y4[-1]), xytext=(x4[-1]-35, y4[-1]+0.04), arrowprops=dict(facecolor='#874C4A', arrowstyle='->'))
    plt.axhline(y=1.0, color='#364B44', linestyle='-', label='Local', linewidth=5)
    plt.axhline(y=0.82, color='#767496', linestyle='--', label='Edge', linewidth=5)
    plt.xlabel('Iterations (%)')
    plt.ylabel('Shrink ratio')
    plt.xlim(left=0)
    plt.legend()
    plt.tight_layout()
    plt.grid(True, color="#999", linestyle='--')
    plt.savefig(os.path.join(os.path.expanduser("~/Desktop"), "ex_05.png"))


def plot_wait_time():
    data = pd.read_csv(os.path.join(data_dir, 'data_02_07.csv'), header=None).values
    plt.figure(figsize=(9, 8), dpi=300)
    x = data[:, 0]
    y1 = data[:, 1]
    y2 = data[:, 2]
    # y3 = data[:, 3]
    y4 = data[:, 4]
    plt.plot(x, y1, marker='^', markerfacecolor='none', markersize=16, markeredgewidth=3, linestyle=':', color='#2F4858', label='CD-Greedy', linewidth=5)
    plt.plot(x, y2, marker='x', markerfacecolor='none', markersize=16, markeredgewidth=3, linestyle=':', color='#19B4A8', label='CD-OPT', linewidth=5)
    plt.plot(x, y4, marker='s', markerfacecolor='none', markersize=16, markeredgewidth=3, linestyle=':', color='#874C4A', label='Random', linewidth=5)
    plt.xlabel('Hover time (s)')
    plt.ylabel('Shrink ratio')
    plt.xlim(left=0)
    plt.legend()
    plt.tight_layout()
    plt.grid(True, color="#999", linestyle='--')
    plt.savefig(os.path.join(os.path.expanduser("~/Desktop"), "ex_06.png"))


def plot_action():
    data = pd.read_csv(os.path.join(data_dir, 'data_02_04.csv'), header=None).values
    plt.figure(figsize=(16, 8), dpi=300)
    x = data[:, 0]
    y1 = data[:, 1]
    y2 = data[:, 2]
    y3 = data[:, 3]
    y4 = data[:, 4]
    y_ = [14, 30, 18, 8, 17, 3, 10]
    a = np.array([0, 1, 2, 3, 4, 5, 6])
    bar_width = 0.3
    colors = ['#2F4858', '#46D0AA', '#1B7C8C']
    plt.bar(a, y1, color=colors[0], width=bar_width, alpha=1, label='CD-Greedy')
    plt.bar(a + bar_width, y2, color=colors[1], width=bar_width, alpha=1, label='CD-OPT')
    plt.bar(a + bar_width * 2, y4, color=colors[2], width=bar_width, alpha=1, label='Random')
    plt.bar(a, y_, color=colors[0], width=bar_width, alpha=0.3)
    plt.bar(a + bar_width, y_, color=colors[1], width=bar_width, alpha=0.3)
    plt.bar(a + bar_width + bar_width, y_, color=colors[2], width=bar_width, alpha=0.3)
    plt.xlabel('UAV trajectory (PT)')
    plt.ylabel('Number of offloading UDs')
    plt.xticks(ticks=[0.3, 1.3, 2.3, 3.3, 4.3, 5.3, 6.3], labels=x.tolist())
    plt.grid(True, axis='y', linestyle=':', color='#333')
    plt.legend()
    plt.tight_layout()
    plt.grid(True, color="#999", linestyle=':')
    plt.savefig(os.path.join(os.path.expanduser("~/Desktop"), "ex_07.png"))


def plot_t_ratio():
    data1 = pd.read_csv(os.path.join(data_dir, 'data_03_01.csv')).values
    data2 = pd.read_csv(os.path.join(data_dir, 'data_03_02.csv')).values
    data3 = pd.read_csv(os.path.join(data_dir, 'data_03_03.csv')).values
    plt.figure(figsize=(9, 8), dpi=300)
    y1 = data1[:, 2][1::5]
    y2 = data2[:, 2][1::5]
    y3 = data3[:, 3]
    x1 = [i * 5 for i in range(len(y1))]
    plt.plot(x1, y1, marker='^', markerfacecolor='none', markersize=16, markeredgewidth=3, linestyle=':', color='#2F4858', label='CD-Greedy', linewidth=5)
    plt.plot(x1, y2, marker='x', markerfacecolor='none', markersize=16, markeredgewidth=3, linestyle=':', color='#19B4A8', label='CD-OPT', linewidth=5)
    x3 = [i * 10 for i in range(len(y3))]
    plt.plot(x3, y3, marker='s', markerfacecolor='none', markersize=16, markeredgewidth=3, linestyle=':', color='#874C4A', label='Random', linewidth=5)
    plt.xlabel('Episode (%)')
    plt.ylabel('Time utilization')
    plt.xlim(left=0)
    plt.legend()
    plt.tight_layout()
    plt.grid(True, color="#999", linestyle='--')
    plt.savefig(os.path.join(os.path.expanduser("~/Desktop"), "ex_08.png"))


def plot_cpu_ratio():
    data1 = pd.read_csv(os.path.join(data_dir, 'data_03_01.csv')).values
    data2 = pd.read_csv(os.path.join(data_dir, 'data_03_02.csv')).values
    data3 = pd.read_csv(os.path.join(data_dir, 'data_03_03.csv')).values
    plt.figure(figsize=(9, 8), dpi=300)
    y1 = data1[:, 3][1::5]
    y2 = data2[:, 3][1::5]
    y3 = data3[:, 4]
    x1 = [i * 5 for i in range(len(y1))]
    plt.plot(x1, y1, marker='^', markerfacecolor='none', markersize=16, markeredgewidth=3, linestyle=':', color='#2F4858', label='CD-Greedy', linewidth=5)
    plt.plot(x1, y2, marker='x', markerfacecolor='none', markersize=16, markeredgewidth=3, linestyle=':', color='#19B4A8', label='CD-OPT', linewidth=5)
    x3 = [i * 10 for i in range(len(y3))]
    plt.plot(x3, y3, marker='s', markerfacecolor='none', markersize=16, markeredgewidth=3, linestyle=':', color='#874C4A', label='Random', linewidth=5)
    plt.xlabel('Episode (%)')
    plt.ylabel('CPU utilization')
    plt.xlim(left=0)
    plt.legend()
    plt.tight_layout()
    plt.grid(True, color="#999", linestyle='--')
    plt.savefig(os.path.join(os.path.expanduser("~/Desktop"), "ex_09.png"))


def plot_scale():
    data = pd.read_csv(os.path.join(data_dir, 'data_04_01.csv')).values
    plt.figure(figsize=(16, 8), dpi=300)
    x = data[:, 1]
    y1 = data[:, 2]
    bar_width = 30
    colors = ['#2F4858', '#46D0AA', '#1B7C8C']
    # plt.bar(x, y1, color=colors[0], width=bar_width, alpha=1, label='CD-Greedy')
    plt.plot(x, y1, color=colors[0], alpha=1, label='CD-Greedy')
    # plt.plot(x, y1 + 1, marker='x', markersize=12, linestyle='-', color=colors[0], linewidth=3)
    plt.xlabel('Number of UDs')
    plt.ylabel('Number of iterations for convergence')
    plt.grid(True, linestyle=':', color='#333')
    plt.tight_layout()
    plt.savefig(os.path.join(os.path.expanduser("~/Desktop"), "ex_10.png"))


def plot_scale_time():
    colors = ['#2F4858', '#46D0AA', '#874C4A']
    data = pd.read_csv(os.path.join(data_dir, 'data_04_02.csv'), header=0).values
    plt.figure(figsize=(9, 8), dpi=300)
    x1 = np.array(data[:10, 1], dtype=int)
    x2 = np.array(data[10:20, 1], dtype=int)
    x3 = np.array(data[20:30, 1], dtype=int)
    t1 = np.array(data[:10, 6], dtype=float)
    t2 = np.array(data[10:20, 6], dtype=float)
    t3 = np.array(data[20:30, 6], dtype=float)
    plt.plot(x1, t1, marker='^', markerfacecolor='none', markersize=16, markeredgewidth=2, linestyle='-', color=colors[0], label='CD-Greedy', linewidth=5)
    plt.plot(x2, t2, marker='x', markerfacecolor='none', markersize=16, markeredgewidth=2, linestyle='-', color=colors[1], label='CD-OPT', linewidth=5)
    plt.plot(x3, t3, marker='s', markerfacecolor='none', markersize=16, markeredgewidth=2, linestyle='-', color=colors[2], label='Random', linewidth=5)
    plt.xlabel('Number of IoTDs')
    plt.ylabel('Time per iteration (s)')
    plt.xlim(left=50)
    plt.yticks([0, 10, 20, 30, 40, 50, 60, 70, 80], ['0', '10', '20', '30', '40', '50', '60', '70', '∞'])
    plt.legend()
    plt.tight_layout()
    plt.grid(True, color="#999", linestyle='--')
    plt.savefig(os.path.join(os.path.expanduser("~/Desktop"), "ex_11.png"))


def plot_scale_sr():
    colors = ['#2F4858', '#46D0AA', '#874C4A']
    data = pd.read_csv(os.path.join(data_dir, 'data_04_02.csv'), header=0).values
    plt.figure(figsize=(9, 8), dpi=300)
    x1 = np.array(data[:10, 1], dtype=int)
    x2 = np.array(data[10:20, 1], dtype=int)
    x3 = np.array(data[20:30, 1], dtype=int)
    sr1 = np.array(data[:10, 4], dtype=float)
    sr2 = np.array(data[10:20, 4], dtype=float)
    sr3 = np.array(data[20:30, 4], dtype=float)
    plt.plot(x1, sr1/x1, marker='^', markerfacecolor='none', markersize=16, markeredgewidth=3, linestyle='-', color=colors[0], label='CD-Greedy', linewidth=5)
    plt.plot(x2, sr2/x2, marker='x', markerfacecolor='none', markersize=16, markeredgewidth=3, linestyle='-', color=colors[1], label='CD-OPT', linewidth=5)
    plt.plot(x3, sr3/x3, marker='s', markerfacecolor='none', markersize=16, markeredgewidth=3, linestyle='-', color=colors[2], label='Random', linewidth=5)
    plt.xlabel('Number of IoTDs')
    plt.ylabel('Shrinkage ratio')
    plt.xlim(left=50)
    plt.legend()
    plt.tight_layout()
    plt.grid(True, color="#999", linestyle='--')
    plt.savefig(os.path.join(os.path.expanduser("~/Desktop"), "ex_12.png"))


plot_scale_time()
